# ------------------------------------------------------------------------------
# Gets scores from lasso
# Author: Cassidy Shubatt <cshubatt@gmail.com>
# To run: bsub -q medium -R "rusage[mem=1000]" bash 03_predict_lasso.sh {outcome}
# ------------------------------------------------------------------------------

# Seeding ----------------------------------------------------------------------
set.seed(1)

# Libraries --------------------------------------------------------------------
message("Loading libraries...")
library(here)
library(yaml)
library(optparse)
library(tidyverse)
library(glue)
library(Matrix)
library(glmnet)

temp <- here(
  "code", "06_physician_boundedness", "01_behavioral_lasso", "temp"
)
u <- modules::use(here("lib", "util.R"))

# Command Line -----------------------------------------------------------------
option_list <- list(
  make_option("--target", type = "character", help = "name of target variable")
)
args <- parse_args(OptionParser(option_list = option_list))

# Load Data --------------------------------------------------------------------
message("Loading data...")
paths <- read_yaml(here("lib", "filepaths.yml"))
overnight_lab <- ""
split <- "random"
cohort_fp <- file.path(paths$modeling$dir, "cohorts", "random")
x <- readRDS(glue(paths$features$test))
ids <- readRDS(file.path(cohort_fp, "test_cohort.rds"))

model <- readRDS(file.path(temp, glue("lasso__{args$target}.rds")))

# Predict ----------------------------------------------------------------------
message("Making predictions...")
predictions <- predict(model, x, s = model$lambda, type = "response")

message("Formatting predictions...")
prediction_tb <- tibble(
  lambda = model$lambda,
  n_coef = colSums(coef(model) != 0) - 1,
  order_name = args$target,
  model_name = "lasso",
  score_name = args$target,
  score = map(
    seq_along(model$lambda),
    ~ as.numeric(predictions[, .x])
  )
)
prediction_tb <- prediction_tb %>%
  group_by(n_coef) %>%
  top_n(1, -lambda) %>%
  ungroup()

# Save -------------------------------------------------------------------------
message("Saving...")
write_rds(prediction_tb, file.path(temp, glue("scores__lasso__{args$target}.rds")))

# ------------------------------------------------------------------------------
message("Done.")
